import os
import mysql.connector

def init_database():
    """初始化数据库和表格"""
    # 数据库配置
    db_config = {
        'host': os.getenv('MYSQL_HOST', 'localhost'),
        'user': os.getenv('MYSQL_USER', 'root'),
        'password': os.getenv('MYSQL_PASSWORD', 'neo..6555894')
    }
    
    try:
        # 连接MySQL服务器（不指定数据库）
        conn = mysql.connector.connect(
            host=db_config['host'],
            user=db_config['user'],
            password=db_config['password']
        )
        cursor = conn.cursor()
        
        # 创建数据库（如果不存在）
        db_name = os.getenv('MYSQL_DATABASE', 'sales_chat_mvp')
        cursor.execute(f"CREATE DATABASE IF NOT EXISTS {db_name} CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci")
        print(f"数据库 '{db_name}' 已创建或已存在")
        
        # 切换到新创建的数据库
        cursor.execute(f"USE {db_name}")
        
        # 创建表（如果不存在）
        create_table_sql = """
        CREATE TABLE IF NOT EXISTS chat_history (
            id INT AUTO_INCREMENT PRIMARY KEY,
            nickname VARCHAR(255) COMMENT 'WeChat nickname or group name',
            is_group BOOLEAN COMMENT 'Whether it is a group chat',
            customer_nickname VARCHAR(255) COMMENT 'Customer WeChat nickname',
            chat_history TEXT COMMENT 'Chat history in JSON format',
            ai_analysis TEXT COMMENT 'AI analysis record',
            analysis_time DATETIME COMMENT 'Time when AI analysis was performed',
            ai_analysis_time DATETIME COMMENT 'Time when AI analysis was updated',
            week_start_time DATETIME COMMENT 'Start time of the week',
            end_time DATETIME COMMENT 'Actual end time of the chat',
            time_span TEXT COMMENT 'Time span in JSON format',
            ai_judgment JSON COMMENT 'First part - AI judgment of chat relevance and customer nickname',
            ai_sales_stage JSON COMMENT 'Second part - Sales stage identification',
            ai_customer_profile JSON COMMENT 'Third part - Customer profile analysis',
            ai_deal_intention JSON COMMENT 'Fourth part - Deal intention evaluation',
            ai_communication_quality JSON COMMENT 'Fifth part - Communication quality scoring',
            ai_tags JSON COMMENT 'Sixth part - Generated tags'
        ) CHARACTER SET utf8mb4 COLLATE utf8mb4_unicode_ci;
        """
        cursor.execute(create_table_sql)
        print("表 'chat_history' 已创建或已存在")
        
        # 检查是否需要添加analysis_time字段
        check_column_sql = """
        SELECT COUNT(*)
        FROM information_schema.COLUMNS
        WHERE TABLE_SCHEMA = %s
        AND TABLE_NAME = 'chat_history'
        AND COLUMN_NAME = 'analysis_time'
        """
        cursor.execute(check_column_sql, (db_name,))
        column_exists = cursor.fetchone()[0]
        
        # 如果字段不存在，则添加
        if column_exists == 0:
            add_column_sql = """
            ALTER TABLE chat_history
            ADD COLUMN analysis_time DATETIME COMMENT 'Time when AI analysis was performed'
            """
            cursor.execute(add_column_sql)
            print("已添加 'analysis_time' 字段到 'chat_history' 表")
            
        # 检查是否需要添加ai_analysis_time字段
        check_column_sql = """
        SELECT COUNT(*)
        FROM information_schema.COLUMNS
        WHERE TABLE_SCHEMA = %s
        AND TABLE_NAME = 'chat_history'
        AND COLUMN_NAME = 'ai_analysis_time'
        """
        cursor.execute(check_column_sql, (db_name,))
        column_exists = cursor.fetchone()[0]
        
        # 如果字段不存在，则添加
        if column_exists == 0:
            add_column_sql = """
            ALTER TABLE chat_history
            ADD COLUMN ai_analysis_time DATETIME COMMENT 'Time when AI analysis was updated'
            """
            cursor.execute(add_column_sql)
            print("已添加 'ai_analysis_time' 字段到 'chat_history' 表")
        
        # 检查是否需要添加end_time字段
        check_column_sql = """
        SELECT COUNT(*)
        FROM information_schema.COLUMNS
        WHERE TABLE_SCHEMA = %s
        AND TABLE_NAME = 'chat_history'
        AND COLUMN_NAME = 'end_time'
        """
        cursor.execute(check_column_sql, (db_name,))
        column_exists = cursor.fetchone()[0]
        
        # 如果字段不存在，则添加
        if column_exists == 0:
            add_column_sql = """
            ALTER TABLE chat_history
            ADD COLUMN end_time DATETIME COMMENT 'Actual end time of the chat'
            """
            cursor.execute(add_column_sql)
            print("已添加 'end_time' 字段到 'chat_history' 表")
        
        # 检查并添加6个新字段
        new_columns = [
            ('ai_judgment', 'JSON COMMENT "First part - AI judgment of chat relevance and customer nickname"'),
            ('ai_sales_stage', 'JSON COMMENT "Second part - Sales stage identification"'),
            ('ai_customer_profile', 'JSON COMMENT "Third part - Customer profile analysis"'),
            ('ai_deal_intention', 'JSON COMMENT "Fourth part - Deal intention evaluation"'),
            ('ai_communication_quality', 'JSON COMMENT "Fifth part - Communication quality scoring"'),
            ('ai_tags', 'JSON COMMENT "Sixth part - Generated tags"')
        ]
        
        for column_name, column_def in new_columns:
            check_column_sql = """
            SELECT COUNT(*)
            FROM information_schema.COLUMNS
            WHERE TABLE_SCHEMA = %s
            AND TABLE_NAME = 'chat_history'
            AND COLUMN_NAME = %s
            """
            cursor.execute(check_column_sql, (db_name, column_name))
            column_exists = cursor.fetchone()[0]
            
            if column_exists == 0:
                add_column_sql = f"""
                ALTER TABLE chat_history
                ADD COLUMN {column_name} {column_def}
                """
                cursor.execute(add_column_sql)
                print(f"已添加 '{column_name}' 字段到 'chat_history' 表")
        
        # 提交事务
        conn.commit()
        print("数据库初始化完成")
    
    except Exception as e:
        print(f"初始化数据库时出错: {e}")
    
    finally:
        if 'conn' in locals() and conn.is_connected():
            cursor.close()
            conn.close()

if __name__ == "__main__":
    init_database() 